import cv2
import threading, time
import numpy as np
from xgolib import XGO
from PIL import Image,ImageDraw,ImageFont
import xgoscreen.LCD_2inch as LCD_2inch
import os
from ultralytics import YOLO

class DogController():
    def __init__(self):
        self.dog = XGO("xgolite")
        self.display = LCD_2inch.LCD_2inch()
        self.display.Init() 
        self.display.clear()
        self.splash = Image.new("RGB",(320,240),"black")
        self.display.ShowImage(self.splash)
        self.cap = None
        self.camera_still = False
        self.save_path = "./saved_images"
        # 确保保存路径存在
        if not os.path.exists(self.save_path):
            os.makedirs(self.save_path)
        self.model = YOLO("best.pt")
        
        # 优化参数
        self.frame_skip = 2  # 每2帧处理一次检测，减少计算负担
        self.current_frame_count = 0
        # 定义颜色映射
        self.color_mapping = {
            # 长方条类别
            'bluerec': "blue",
            'greenrec': "green",
            'redrec': "red",
            
            # 圆柱台类别
            'bluecylinder': "blue",
            'greencylinder': "green",
            'redcylinder': "red",
        }
    def open_camera(self):
        print("📷 正在打开摄像头...")
        if self.cap is None:
            self.cap = cv2.VideoCapture(0)
            # self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
            # self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
            # 关键优化：设置缓冲区大小为1，减少延迟
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
            # 提高帧率
            self.cap.set(cv2.CAP_PROP_FPS, 30)
            print("✅ 摄像头初始化完成 (320x240@30fps, 低延迟模式)")
    
    def save_current_image(self):
        """保存当前帧图像"""
        if self.cap is not None and self.cap.isOpened():
            success, image = self.cap.read()
            if success:
                timestamp = time.strftime("%Y%m%d_%H%M%S")
                filename = f"dog_capture_{timestamp}.jpg"
                filepath = os.path.join(self.save_path, filename)
                cv2.imwrite(filepath, image)
                print(f"📸 图像已保存到: {filepath}")
            else:
                print("⚠️ 无法获取当前帧")
        else:
            print("⚠️ 摄像头未开启")
    def camera_mode(self):
        if not self.cap.isOpened():
            print("无法打开摄像头")
            return
        print("🎯 实时检测模式启动（优化版）")
        while self.camera_still and self.cap is not None and self.cap.isOpened():
            # 清空缓冲区，获取最新帧
            ret = True
            frame = None
            # 连续读取几帧，丢弃旧帧，只保留最新的
            for _ in range(3):  # 读取3帧，只用最后一帧
                ret, frame = self.cap.read()
                if not ret:
                    break
            if not ret or frame is None:
                print("读取摄像头失败")
                time.sleep(0.01)
                continue
            # self.current_frame_count += 1
            # # 每隔frame_skip帧才进行一次检测，减少计算负担
            # if self.current_frame_count % self.frame_skip == 0:
            #     # YOLO检测 - 降低置信度阈值，提高检测灵敏度
            #     results = self.model.predict(frame, conf=0.2, verbose=False, imgsz=320)
            #     # 绘制检测结果
            #     for r in results:
            #         if r.boxes is not None and len(r.boxes) > 0:
            #             for box in r.boxes:
            #                 x1, y1, x2, y2 = map(int, box.xyxy[0])
            #                 conf = float(box.conf[0])
            #                 cls_id = int(box.cls[0])
            #                 label = self.model.names[cls_id]
            #                 text = f"{label} {conf:.2f}"
            #                 print(f"label:{label} - x1:{x1},y1:{y1}")
            #                 # 绘制矩形框和标签
            #                 cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                            # cv2.putText(frame, text, (x1, y1 - 10), 
                            #            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            # 无论是否检测，都显示当前帧（保证流畅度）
            try:
                # 色彩通道转换 BGR -> RGB
                image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                image_rgb = cv2.flip(image_rgb, 1)  # 水平翻转
                pil_image = Image.fromarray(image_rgb)
                # 显示处理后的图像
                self.display.ShowImage(pil_image)
            except Exception as e:
                print(f"显示图像时出错: {e}")
                
            # 减少延迟 - 更小的sleep时间
            time.sleep(0.01)  # 约100fps显示，但检测频率较低

    def xgoCamera(self, switch):
        print(f"🎮 摄像头开关: {'开启' if switch else '关闭'}")
        if switch:
            self.open_camera()
            if self.cap and self.cap.isOpened():
                self.camera_still = True
                # 重置帧计数器
                self.current_frame_count = 0
                
                # 启动摄像头线程
                t = threading.Thread(target=self.camera_mode, daemon=True)  
                t.start() 
                # 调整摄像头角度
                print("🔧 调整机械狗摄像头角度...")
                self.dog.motor(11, -30)
                self.dog.motor(21, -30)
                print("✅ 摄像头角度调整完成")
            else:
                print("❌ 摄像头打开失败")
        else:
            print("🛑 正在关闭摄像头...")
            self.camera_still = False
            time.sleep(0.1)  # 给线程时间停止
            if self.cap:
                self.cap.release()
                self.cap = None
            splash = Image.new("RGB", (320, 240), "black")
            self.display.ShowImage(splash)
            print("✅ 摄像头已关闭")

    def close(self):
        """清理资源"""
        print("🧹 正在清理系统资源...")
        self.camera_still = False
        time.sleep(0.2)  # 等待线程停止
        if self.cap:
            self.cap.release()
            print("✅ 摄像头资源已释放")
        self.display.clear()
        print("✅ 显示屏已清理")
        print("💤 系统关闭完成")

# 使用示例
if __name__ == "__main__":
    print("🤖 绝影机械狗圆柱台和长条检测系统 (优化版)")
    print("=" * 60)
    print("📌 优化特性:")
    print("   • 低延迟缓冲区设置")
    print("   • 帧跳跃检测策略")
    print("   • 实时显示优化")
    print("📌 按 Ctrl+C 退出程序")
    print("=" * 60)
    
    x = DogController()
    try:
        x.xgoCamera(True)
        # 运行检测 (无限循环，直到用户中断)
        while True:
            input("按下enter键开始拍照")
            x.save_current_image()
    except KeyboardInterrupt:
        print("\n\n⚠️  用户中断程序")
        print("🛑 正在安全关闭系统...")
        x.xgoCamera(False)
    finally:
        x.close()
        print("\n🎯 程序运行结束")
        print("=" * 60)